# -*- coding: utf-8 -*-
# authors 13077/13130
# data: 30 de Novembro de 2014
#

from pylab import *
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
import time

FIGURE_PATH = '../data/foo.png'


def plot_simple(graph_title, graph_xlabel, graph_ylabel, list_year, list_value, reg, avg, med):

    """
    Plots a line graph based on two lists of values
    :param graph_title: Graph title
    :param graph_xlabel: X axis label
    :param graph_ylabel: Y axis label
    :param list_year: First list of values (x axis)
    :param list_value: Second list of values (y axis)
    :param reg: Plot linear regression
    :param avg: Plot average
    :param med: Plot median
    """
    t1 = time.time()
    start_year = min(list_year)
    end_year = max(list_year)
    plt.figure()
    plt.ylabel(graph_ylabel)
    plt.xlabel(graph_xlabel)
    plt.title(graph_title + ' | ' + str(start_year) + '-' + str(end_year))

    if reg is True:
        fit = polyfit(list_year, list_value, 1)  # cria lista de uma funçao de 1º grau pela interpolação dos valores
        fit_fn = poly1d(fit)  # fit é agora uma função que recebe x e devolve f(x)
        plot(list_year, fit_fn(list_year), '--k', label=u'Tendência', linewidth=1.5)  # linear regression plot

    if avg is True:
        avg = average(list_value)
        avg_label = u'Média: ' + str(avg)
        plot([start_year, end_year], [avg, avg], label=avg_label,
             color='#db12fd', linewidth=2)  # average plot

    if med is True:
        med = median(list_value)
        med_label = u'Mediana: ' + str(med)
        plot([start_year, end_year], [med, med], label=med_label,
             color='#71fb12', linewidth=2)  # median plot

    plt.plot(list_year, list_value, label=graph_title, linewidth=3)  # line
    plt.grid()
    plt.xticks(ticks(list_year), map(lambda x: "%g" % x, ticks(list_year)))
    plt.yticks(ticks(list_value))
    plt.legend(loc=0, ncol=1, borderaxespad=0.)
    savefig(FIGURE_PATH)
    plt.close()
    print "plot_simple exec time: " + str(round(time.time() - t1, 3))


def plot_correlation(list_1, list_2, start, end, xtitle, ytitle):
    """
    Plots a correlation graph based on two lists of values
    :param list_1: First list of values
    :param list_2: Second list of values
    :param start: Starting value
    :param end: Last value
    :param xtitle: X axis label
    :param ytitle: Y axis label
    """
    t1 = time.time()
    plt.figure()
    plot(list_1, list_2, 'yo', color='red')
    fit = polyfit(list_1, list_2, 1)  # cria lista de uma funçao de 1º grau pela interpolação dos valores
    fit_fn = poly1d(fit)  # fit é agora uma função que recebe x e devolve f(x)
    plot(list_1, fit_fn(list_1), '--k', label=u'Tendência', linewidth=1.5)  # linear regression plot

    r = pearsonr(list_1, list_2)
    r1 = round(float(r[0]), 3)
    title(
        u'Correlação entre ' + xtitle + ' e ' + ytitle + '\n' + 'r = ' + str(r1) + ' | ' + str(start) + ' - ' + str(
            end))
    plt.xticks(ticks(list_1))
    plt.yticks(ticks(list_2))
    plt.ylabel(ytitle)
    plt.xlabel(xtitle)
    savefig(FIGURE_PATH)
    plt.close()
    print "plot_correlation exec time: " + str(round(time.time() - t1, 3))


def ticks(nrange):
    """
    Creates a list of ticks for an axis
    :param nrange: List with a range of values
    :return: The list of ticks
    """
    t1 = time.time()
    if len(nrange) < 8:
        return nrange
    else:
        res = list(np.arange(min(nrange), max(nrange), (max(nrange) - min(nrange)) / 4))
    if max(nrange) != max(res):
        res.append(max(nrange))
    print "ticks exec time: " + str(round(time.time() - t1, 3))
    return res


def average(values):
    """
    Calculates the average of a list of values
    :param values: List of values
    :return: The average value
    """
    t1 = time.time()
    a = sum(values) / len(values)
    avg = round(a, 2)
    print "average exec time: " + str(round(time.time() - t1, 3))
    return avg


def median(values):
    """
    Calculates the median of a list of values
    :param values: List of values
    :return: The median value
    """
    length = len(values)
    middle = len(values) / 2

    if len(values) == 2:
        return round(average(values))
    if length % 2 == 0:
        return round((values[middle] + values[middle - 1]) / 2, 2)
    else:
        return round(values[middle], 2)